{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/opensim-org/opensim-core/blob/main/Bindings/Python/tutorials/Tutorial%208%20-%20Creating%20Muscle-driven%20Simulations%20with%20OpenSim%20Moco.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YS-0eq5Y6BjU"
      },
      "source": [
        "# 8. Creating muscle-driven simulations with OpenSim Moco"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZClpIjbitD4"
      },
      "source": [
        "\n",
        "## 8.1. Objectives"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AKANMMa6iycq"
      },
      "source": [
        "**Introduction**\n",
        "\n",
        "[OpenSim Moco](https://opensim-org.github.io/opensim-moco-site/) is a software package for creating musculoskeletal simulations using trajectory optimization [[1]](https://doi.org/10.1371/journal.pcbi.1008493). Users define problems by providing an OpenSim model, cost function terms, and additional problem constraints via an easy-to-use scripting interface. Moco rapidly solves trajectory optimization problems via the state-of-the-art direct collocation method.\n",
        "\n",
        "**Purpose**\n",
        "\n",
        "The purpose of this tutorial is to showcase Moco's advanced simulation capabilities using the [OpenSim conda package](https://anaconda.org/opensim-org/opensim) and [Jupyter notebooks](https://jupyter.org/). In this tutorial you will:\n",
        "\n",
        "*   Set up a new conda environment and install OpenSim with Moco.\n",
        "*   Learn how to replace muscle geometry paths with smooth function-based paths better suited for optimization.\n",
        "*   Pose and solve a muscle-driven optimization problem for simulating a squat-to-stand motion.\n",
        "*   Design a passive assistive device at the knee using parameter optimization.\n",
        "*   Investigate how neural control limitations affect the squat-to-stand motion using muscle synergies.\n",
        "\n",
        "**Resources**\n",
        "\n",
        "All of the files necessary to complete this tutorial are available in the following Google Drive folder: https://drive.google.com/drive/folders/1eGe2UxcsYmT3OXQN3AdpgnrJy803sdci?usp=sharing\n",
        "\n",
        "While you complete the tutorial, you will find cells that automatically download the required files.\n",
        "\n",
        "To obtain a link to one of the files in Google Drive, right-click over it and click on **Get Link**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kPnGlZIT8oiM"
      },
      "source": [
        "## 8.2. Set up conda environment and download resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3kxYrs4Hi6Nl"
      },
      "source": [
        "In the following cell, you will use [condacolab](https://github.com/conda-incubator/condacolab) to install [Miniconda](https://docs.conda.io/en/latest/miniconda.html). If you want to install a different conda distribution (e.g., Anaconda, Mambaforge) check [condacolab's documentation](https://github.com/conda-incubator/condacolab#usage)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5-sX9L3f6Asq",
        "outputId": "07c714c8-7744-4bf4-fa0f-f1c66d2fcec9"
      },
      "outputs": [],
      "source": [
        "!pip install -q condacolab\n",
        "import condacolab\n",
        "condacolab.install()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LdmpfqC8Dt0x"
      },
      "source": [
        "Don't worry if after executing the previous cell you get an error saying that your session has failed. The reason for this is that condacolab needs to restart the session for the changes to take effect. Because of this, you have to execute the previous cell first, before executing any other cell.\n",
        "\n",
        "Now, execute the following cell to install the OpenSim conda package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H9WE2h3fD61J",
        "outputId": "2795df96-b021-40eb-ba47-39a105cc2ac6"
      },
      "outputs": [],
      "source": [
        "!conda install opensim-org::opensim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, execute the following cell to create a button that will download the \"squat-to-stand\" OpenSim [`Model`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1Model.html). This is a single leg with 3 DOF and 9 muscles, and the foot is welded to the ground. The model file and accompanying geometry files will be extracted to the working directory.\n",
        "\n",
        "![model](https://drive.google.com/uc?export=view&id=1YE2lS4_jCnsFsev0GSbR_OGiXpXC3iMS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# @title Download the \"squat-to-stand\" model and geometry files { display-mode: \"form\" }\n",
        "# @markdown After executing this cell, a download button will appear. Click the button to download and extract the zip file containing the \"squat-to-stand\" model and geometry files.\n",
        "\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "import requests\n",
        "import zipfile\n",
        "import io\n",
        "\n",
        "# Define the function to download and extract the zip file\n",
        "def download_and_extract_zip(url):\n",
        "    try:\n",
        "        # Get the content of the URL\n",
        "        response = requests.get(url)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        # Create a ZipFile object from the downloaded content\n",
        "        with zipfile.ZipFile(io.BytesIO(response.content)) as zip_ref:\n",
        "            # Extract all the contents into the working directory\n",
        "            zip_ref.extractall()\n",
        "            print(\"Extraction complete!\")\n",
        "    except requests.exceptions.RequestException as e:\n",
        "        print(f\"Failed to download the file: {e}\")\n",
        "    except zipfile.BadZipFile:\n",
        "        print(\"The file is not a valid zip file.\")\n",
        "    except Exception as e:\n",
        "        print(f\"An error occurred: {e}\")\n",
        "\n",
        "# Create a button to trigger the download and extraction\n",
        "download_button = widgets.Button(\n",
        "    description='Download Model',\n",
        "    disabled=False,\n",
        "    button_style='',\n",
        "    tooltip='Click to download the squat-to-stand model and geometry files',\n",
        "    icon='download'\n",
        ")\n",
        "\n",
        "# Define the button click event handler\n",
        "def on_button_click(b):\n",
        "    url = 'https://drive.google.com/uc?export=download&id=1xnrBNTqIHQVIIxQNQEbHUmVLLRduV5jd'\n",
        "    if url:\n",
        "        download_and_extract_zip(url)\n",
        "    else:\n",
        "        print(\"Please enter a valid URL.\")\n",
        "\n",
        "# Attach the event handler to the button\n",
        "download_button.on_click(on_button_click)\n",
        "\n",
        "# Display the widgets\n",
        "display(download_button)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8.3. Fit function-based paths for the squat-to-stand model \n",
        "\n",
        "Before we create a squat-to-stand simulation, we will replace the muscle geometry in the model with smooth polynomial functions which, being smooth and continuous everywhere, lead to better convergence in gradient-based optimization. The original model represents muscle geometry using [`GeometryPath`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1GeometryPath.html) elements. We will replace these with [`FunctionBasedPath`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1FunctionBasedPath.html) elements, which will compute muscle length, speed, and moment arms using OpenSim's [`MultivariatePolynomialFunction`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1MultivariatePolynomialFunction.html) class. The polynomials need to be multivariate since muscle paths can depend on multiple coordinates in the model (e.g., the gastrocnemius path depends on both the knee and ankle coordinates)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "First, run the following cell to import OpenSim and set the logging file. If successful, it will print version 4.5.1 with build date in July 2024."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import opensim as osim\n",
        "osim.Logger.removeFileSink()\n",
        "osim.Logger.addFileSink('opensim.log')\n",
        "osim.GetVersionAndDate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, run the following cell to use forward integration to create a random trajectory that we can use as a reference when fitting polynomial functions to represent muscle lengths. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Load the model and initialize the multibody system\n",
        "model = osim.Model('squatToStand_3dof9musc.osim')\n",
        "state = model.initSystem()\n",
        "\n",
        "# Set an initial pose.\n",
        "coordSet = model.updCoordinateSet()\n",
        "coordSet.get('hip_flexion_r').setValue(state, -40 * (np.pi / 180))\n",
        "coordSet.get('knee_angle_r').setValue(state, -45 * (np.pi / 180))\n",
        "coordSet.get('ankle_angle_r').setValue(state, -25 * (np.pi / 180))\n",
        "\n",
        "# Run a 0.4 second forward simulation.\n",
        "manager = osim.Manager(model, state)\n",
        "manager.integrate(0.4)\n",
        "\n",
        "# Save the states to a file.\n",
        "table = manager.getStatesTable()\n",
        "sto = osim.STOFileAdapter()\n",
        "if not os.path.exists('path_fitting'):\n",
        "    os.mkdir('path_fitting')\n",
        "sto.write(table, os.path.join('path_fitting', 'reference_states.sto'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will now use [`PolynomialPathFitter`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1PolynomialPathFitter.html), a utility class for automating the path fitting process. First, it will randomly sample coordinate values around the reference trajectory. Then, muscle lengths and moment arms will be computed from the sampled coordinate values. Coefficients for multivariate polynomial functions representing the muscle lengths are determined using a least-squares fit. Both the muscle lengths and the moment arms (via the polynomial derivatives) are included in the fitting process. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Construct a PolynomialPathFitter object.\n",
        "fitter = osim.PolynomialPathFitter()\n",
        "\n",
        "# Set the model.\n",
        "fitter.setModel(osim.ModelProcessor('squatToStand_3dof9musc.osim'))\n",
        "\n",
        "# Set the reference states.\n",
        "table = osim.TimeSeriesTable(os.path.join('path_fitting', 'reference_states.sto'))\n",
        "fitter.setCoordinateValues(osim.TableProcessor(table))\n",
        "\n",
        "# Configure settings.\n",
        "fitter.setMaximumPolynomialOrder(6)\n",
        "fitter.setNumSamplesPerFrame(10)\n",
        "fitter.setGlobalCoordinateSamplingBounds(osim.Vec2(-30, 30))\n",
        "fitter.setUseStepwiseRegression(True)\n",
        "fitter.setPathLengthTolerance(1e-4)\n",
        "fitter.setMomentArmTolerance(1e-4)\n",
        "fitter.setOutputDirectory('path_fitting')\n",
        "\n",
        "# Run the path fitter.\n",
        "fitter.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we'll plot the results to compare the muscle lengths computed from the fitted polynomial functions to the original muscle lengths."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# @title Plot the path fitting results { display-mode: \"form\" }\n",
        "# @markdown Execute this cell to plot the fitted path lengths together with the original path lengths.\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "# Load the original and fitted path lengths and the path length based on the\n",
        "# randomly sampled coordinates values.\n",
        "lengths = osim.TimeSeriesTable(\n",
        "    os.path.join('path_fitting', 'squatToStand_path_lengths.sto'))\n",
        "lengthsFitted = osim.TimeSeriesTable(\n",
        "    os.path.join('path_fitting', 'squatToStand_path_lengths_fitted.sto'))\n",
        "lengthsSampled = osim.TimeSeriesTable(\n",
        "    os.path.join('path_fitting', 'squatToStand_path_lengths_sampled.sto'))\n",
        "lengthsSampledFitted = osim.TimeSeriesTable(\n",
        "    os.path.join('path_fitting', 'squatToStand_path_lengths_sampled_fitted.sto'))\n",
        "\n",
        "# Plot the path lengths. Blue is the original, orange is the fitted.\n",
        "fig, axs = plt.subplots(3, 3, figsize=(8, 8))\n",
        "for ilabel, label in enumerate(lengths.getColumnLabels()):\n",
        "    irow = ilabel // 3\n",
        "    icol = ilabel % 3\n",
        "    axs[irow, icol].scatter(\n",
        "        lengthsSampled.getIndependentColumn(),\n",
        "        lengthsSampled.getDependentColumn(label).to_numpy(),\n",
        "        alpha=0.4, color='blue', s=0.3)\n",
        "    axs[irow, icol].scatter(\n",
        "        lengthsSampledFitted.getIndependentColumn(),\n",
        "        lengthsSampledFitted.getDependentColumn(label).to_numpy(),\n",
        "        alpha=0.4, color='orange', s=0.3)\n",
        "    axs[irow, icol].plot(\n",
        "        lengths.getIndependentColumn(),\n",
        "        lengths.getDependentColumn(label).to_numpy(),\n",
        "        label='original', lw=4)\n",
        "    axs[irow, icol].plot(\n",
        "        lengthsFitted.getIndependentColumn(),\n",
        "        lengthsFitted.getDependentColumn(label).to_numpy(),\n",
        "        label='fitted', lw=3.5)\n",
        "    title = label.replace('/forceset/', '').replace('_length', '')\n",
        "    axs[irow, icol].set_title(title)\n",
        "    if not irow and not icol:\n",
        "        axs[irow, icol].legend()\n",
        "    if not irow == 2:\n",
        "        axs[irow, icol].set_xticklabels([])\n",
        "    else:\n",
        "        axs[irow, icol].set_xlabel('time (s)')\n",
        "    if not icol:\n",
        "        axs[irow, icol].set_ylabel('muscle length (m)')\n",
        "\n",
        "fig.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now that we have a good fit for our function-based muscle paths, we can move on to creating a squat-to-stand simulation with Moco."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8.4. Muscle-driven simulation of a squat-to-stand motion\n",
        "\n",
        "In this section, we will use [`OpenSim Moco`](https://simtk.org/api_docs/opensim/api_docs451/mocomainpage.html) to create a predictive simulation of a motion starting from a squat and ending standing up. This is a **boundary-value problem**, since we need to impose hard constraints at the boundaries of the trajectory (i.e., the initial and final time point). The following image describes the coordinate boundary conditions for the squat and stand positions, and the coordinate value and speed bounds across the trajectory.\n",
        "\n",
        "![squat_to_stand](https://drive.google.com/uc?export=download&id=1ObBouDZnM2zrgnk6rsaCigg3qUr3hiLk)\n",
        "\n",
        "[`MocoStudy`](https://simtk.org/api_docs/opensim/api_docs451/mocostudy.html) is the most general interface for constructing and solving trajectory optimization problems in Moco. First, we will use [`MocoProblem`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1MocoProblem.html) to convert the mathematical description of the boundary-value problem above into a form that Moco understands. Next, we will configure a `MocoSolver`, in this case [`MocoCasADiSolver`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1MocoCasADiSolver.html), to specify [the problem transcription](https://simtk.org/api_docs/opensim/api_docs451/mocotheoryguide.html#dircol) and settings for the optimization solver. Finally, we can use the `MocoStudy` to solve the problem and return a [`MocoSolution`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1MocoSolution.html), which we can use to analyze and visualize the result. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Part 1: Create a new MocoStudy.\n",
        "study = osim.MocoStudy()\n",
        "study.setName('squat_to_stand_prediction')\n",
        "\n",
        "# Part 2: Initialize the problem.\n",
        "problem = study.updProblem()\n",
        "\n",
        "# Part 3: Construct a model using the function-based paths. To simplify\n",
        "# the problem, we will ignore tendon compliance, activation dynamics,\n",
        "# and passive fiber forces.\n",
        "modelProcessor = osim.ModelProcessor('squatToStand_3dof9musc.osim')\n",
        "modelProcessor.append(osim.ModOpScaleMaxIsometricForce(2.0))\n",
        "modelProcessor.append(osim.ModOpReplaceMusclesWithDeGrooteFregly2016())\n",
        "modelProcessor.append(osim.ModOpIgnorePassiveFiberForcesDGF())\n",
        "modelProcessor.append(osim.ModOpIgnoreTendonCompliance())\n",
        "modelProcessor.append(osim.ModOpIgnoreActivationDynamics())\n",
        "modelProcessor.append(osim.ModOpScaleActiveFiberForceCurveWidthDGF(1.5))\n",
        "modelProcessor.append(osim.ModOpReplacePathsWithFunctionBasedPaths(\n",
        "    os.path.join('path_fitting', 'squatToStand_FunctionBasedPathSet.xml')))\n",
        "model = modelProcessor.process()\n",
        "problem.setModel(model)\n",
        "\n",
        "# Part 4: Set bounds on the problem.\n",
        "#\n",
        "# problem.setTimeBounds(initial_bounds, final_bounds)\n",
        "# problem.setStateInfo(path, trajectory_bounds, inital_bounds, final_bounds)\n",
        "#\n",
        "# All *_bounds arguments can be set to a range, [lower upper], or to a\n",
        "# single value (equal lower and upper bounds). Empty brackets, [], indicate\n",
        "# using default bounds (if they exist). You may set multiple state infos at\n",
        "# once using setStateInfoPattern():\n",
        "#\n",
        "# problem.setStateInfoPattern(pattern, trajectory_bounds, inital_bounds, ...\n",
        "#       final_bounds)\n",
        "#\n",
        "# This function supports regular expressions in the 'pattern' argument;\n",
        "# use '.*' to match any substring of the state/control path\n",
        "# For example, the following will set all coordinate value state infos:\n",
        "#\n",
        "# problem.setStateInfoPattern('/path/to/states/.*/value', ...)\n",
        "\n",
        "# Time bounds\n",
        "problem.setTimeBounds(0, 1)\n",
        "\n",
        "# Position bounds: the model should start in a squat and finish\n",
        "# standing up.\n",
        "problem.setStateInfo('/jointset/hip_r/hip_flexion_r/value',\n",
        "    [-2, 0.5], -2, 0)\n",
        "problem.setStateInfo('/jointset/knee_r/knee_angle_r/value',\n",
        "    [-2, 0], -2, 0)\n",
        "problem.setStateInfo('/jointset/ankle_r/ankle_angle_r/value',\n",
        "    [-0.75, 0.75], -0.5, 0)\n",
        "\n",
        "# Velocity bounds: all model coordinates should start and end at rest.\n",
        "problem.setStateInfoPattern('/jointset/.*/speed', [], 0, 0)\n",
        "\n",
        "# Part 5: Add a MocoControlCost to the problem.\n",
        "problem.addGoal(osim.MocoControlGoal('control_effort'))\n",
        "\n",
        "# Part 6: Configure the solver.\n",
        "solver = study.initCasADiSolver()\n",
        "solver.set_num_mesh_intervals(25)\n",
        "solver.set_optim_finite_difference_scheme('forward')\n",
        "solver.set_optim_convergence_tolerance(1e-4)\n",
        "solver.set_optim_constraint_tolerance(1e-4)\n",
        "\n",
        "# Part 7: Print the study to a file.\n",
        "if not os.path.exists('squat_to_stand'):\n",
        "    os.mkdir('squat_to_stand')\n",
        "\n",
        "xml_path = os.path.join('squat_to_stand', 'squat_to_stand.xml')\n",
        "if not os.path.exists(xml_path):\n",
        "    study.printToXML(xml_path)\n",
        "\n",
        "# Part 8: Solve! Write the solution to file.\n",
        "predictSolution = study.solve()\n",
        "solutionFile = os.path.join('squat_to_stand', 'squat_to_stand_prediction_solution.sto')\n",
        "predictSolution.write(solutionFile)\n",
        "# study.visualize(predictSolution)\n",
        "\n",
        "# Part 9: Create a PDF report of the solution.\n",
        "output = os.path.join('squat_to_stand', 'squat_to_stand_prediction_report.pdf')\n",
        "report = osim.report.Report(model,\n",
        "                            solutionFile,\n",
        "                            output=output,\n",
        "                            bilateral=True)\n",
        "report.generate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Use the following exercises to modify the example above and learn more about Moco's capabilities.\n",
        "\n",
        "**<u>Exercise 1.</u> Modify the problem to start from a standing position and end in a squat.**\n",
        "\n",
        "**<u>Exercise 2.</u> Change the cost function by adding a new [`MocoGoal`](https://simtk.org/api_docs/opensim/api_docs451/group__mocogoal.html) to the problem.**\n",
        "\n",
        "**<u>Exercise 3.</u> Advanced: create a new problem using [`MocoStateTrackingGoal`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1MocoStateTrackingGoal.html) to track a squat motion data recorded from OpenCap in Demo 1. See [`exampleMocoTrack`](https://simtk.org/api_docs/opensim/api_docs451/exampleMocoTrack_8py-example.html) for an example on how to use this `MocoGoal`.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8.5. Designing a passive knee assistive device\n",
        "\n",
        "Next, we'll modify the squat-to-stand optimization problem to design a passive assistive device at the knee. We will model the \n",
        "device using an OpenSim [`SpringGeneralizedForce`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1SpringGeneralizedForce.html) which applies a knee torque based on linear spring that is a function of the knee joint angle. The spring torque is parameterized by stiffness and rest length parameters (viscosity is set to zero). Note the rest length has angular units (radians), but we use \"rest length\" instead of \"rest angle\" to stay consistent with the OpenSim API documentation.\n",
        "\n",
        "![device_optimization](https://drive.google.com/uc?export=download&id=1LBG1lXRpfIjNGzTNt6xnt4rP-SO04orT)\n",
        "\n",
        "Execute the cell below to load an interactive menu for designing the assistive device. The menu provides two design models: \"Manual\" and \"Optimization\". In \"Manual\" mode, you will select a stiffness between [0, 400] N-m/rad and a rest length between [-1, 1] rad. In \"Optimization\" mode, the problem uses parameter optimization (via [`MocoParameter`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1MocoParameter.html)) to determine stiffness and rest length parameter values.\n",
        "\n",
        "While running the device optimizations, consider the following questions.\n",
        "\n",
        "**<u>Question 1.</u> How much can you reduce muscle effort via manual parameter design?**\n",
        "\n",
        "**<u>Question 2.</u> Do the optimized parameters make sense? Explain.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# @title Design a passive knee assistive device { display-mode: \"form\" }\n",
        "# @markdown Execute the cell to load the interactive device design menu.\n",
        "\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "\n",
        "def run_device_optimization_squat_to_stand(optimize, stiffness=50, rest_length=0):\n",
        "\n",
        "    # Part 1: Load the base squat-to-stand MocoStudy.\n",
        "    study = osim.MocoStudy(os.path.join('squat_to_stand', 'squat_to_stand.xml'))\n",
        "    problem = study.updProblem()\n",
        "    solver = osim.MocoCasADiSolver.safeDownCast(study.updSolver())\n",
        "\n",
        "    # Part 2: Recreate the model, now adding a passive assistive device at the knee\n",
        "    modelProcessor = osim.ModelProcessor('squatToStand_3dof9musc.osim')\n",
        "    modelProcessor.append(osim.ModOpScaleMaxIsometricForce(2.0))\n",
        "    modelProcessor.append(osim.ModOpReplaceMusclesWithDeGrooteFregly2016())\n",
        "    modelProcessor.append(osim.ModOpIgnorePassiveFiberForcesDGF())\n",
        "    modelProcessor.append(osim.ModOpIgnoreTendonCompliance())\n",
        "    modelProcessor.append(osim.ModOpScaleActiveFiberForceCurveWidthDGF(1.5))\n",
        "    modelProcessor.append(osim.ModOpReplacePathsWithFunctionBasedPaths(\n",
        "        os.path.join('path_fitting', 'squatToStand_FunctionBasedPathSet.xml')))\n",
        "    model = modelProcessor.process()\n",
        "\n",
        "    # The device is a SpringGeneralizedForce acting on the knee angle.\n",
        "    device = osim.SpringGeneralizedForce('knee_angle_r')\n",
        "    device.setName('passive_knee_device')\n",
        "    device.setStiffness(stiffness)\n",
        "    device.setRestLength(rest_length)\n",
        "    device.setViscosity(0)\n",
        "    model.addForce(device)\n",
        "\n",
        "    # Part 3: Update the problem with the new model.\n",
        "    problem.setModel(model)\n",
        "\n",
        "    # Part 4: Add MocoParameters to the problem to optimize device stiffness and rest\n",
        "    # length.\n",
        "    if optimize:\n",
        "        stiffness = osim.MocoParameter('stiffness',\n",
        "                '/forceset/passive_knee_device', 'stiffness', osim.MocoBounds(0, 400))\n",
        "        problem.addParameter(stiffness)\n",
        "\n",
        "        rest_length = osim.MocoParameter('rest_length',\n",
        "                '/forceset/passive_knee_device', 'rest_length', osim.MocoBounds(-1, 1))\n",
        "        problem.addParameter(rest_length)\n",
        "\n",
        "    # Part 5: Update the solver with the new problem (i.e., new model and parameters).\n",
        "    solver.resetProblem(problem)\n",
        "    solver.set_optim_max_iterations(500)\n",
        "\n",
        "    # Part 6: Set solver parameter optimization flag.\n",
        "    # -----------------------------------------------\n",
        "    # This flag tells MocoCasADiSolver that initSystem() calls on the model are\n",
        "    # unnecessary to optimize the parameters. Other parameters (e.g., body mass)\n",
        "    # require initSystem(), since they change the model's underlying multibody\n",
        "    # system in Simbody. Since changing parameters in SpringGeneralizedForce does\n",
        "    # not change the multibody system, we can set this flag to false to speed up\n",
        "    # the optimization.\n",
        "    solver.set_parameters_require_initsystem(False)\n",
        "\n",
        "    # Part 7: Solve the problem.\n",
        "    deviceSolutionSealed = study.solve()\n",
        "    deviceSolution = deviceSolutionSealed.unseal()\n",
        "    if not os.path.exists('device_optimization'):\n",
        "        os.mkdir('device_optimization')\n",
        "    solutionFile = os.path.join('device_optimization',\n",
        "        'squat_to_stand_device_optimization_solution.sto')\n",
        "    deviceSolution.write(solutionFile)\n",
        "    # study.visualize(deviceSolution)\n",
        "\n",
        "    # Part 8: Print solution information.\n",
        "    if optimize:\n",
        "        parameters = deviceSolution.getParameters().to_numpy()\n",
        "        print(f'Optimized stiffness: {parameters[0]:.2f} Nm/rad')\n",
        "        print(f'Optimized rest length: {parameters[1]:.2f} rad')\n",
        "    else:\n",
        "        print(f'Using stiffness: {stiffness:.2f} Nm/rad')\n",
        "        print(f'Using rest length: {rest_length:.2f} rad')\n",
        "    print(f'Muscle effort: {deviceSolution.getObjective():.3f}\\n')\n",
        "\n",
        "    # Part 9: Create a PDF report of the solution.\n",
        "    output = os.path.join('device_optimization',\n",
        "            'squat_to_stand_device_optimization_report.pdf')\n",
        "    ref_files = []\n",
        "    ref_files.append(os.path.join('squat_to_stand',\n",
        "            'squat_to_stand_prediction_solution.sto'))\n",
        "    report = osim.report.Report(model,\n",
        "                                solutionFile,\n",
        "                                output=output,\n",
        "                                ref_files=ref_files,\n",
        "                                bilateral=True)\n",
        "    report.generate()\n",
        "\n",
        "\n",
        "def on_submit(b):\n",
        "    if mode_selector.value == 'Manual':\n",
        "        print('Testing manual design...')\n",
        "        run_device_optimization_squat_to_stand(False,\n",
        "                stiffness_slider.value, rest_length_slider.value)\n",
        "    else:\n",
        "        print('Optimizing design...')\n",
        "        run_device_optimization_squat_to_stand(True)\n",
        "\n",
        "# Define layout for consistent width and alignment\n",
        "common_layout = widgets.Layout(width='400px')\n",
        "\n",
        "# Create mode selector dropdown with the common layout\n",
        "mode_selector = widgets.Dropdown(\n",
        "    options=['Manual', 'Optimization'],\n",
        "    value='Manual',\n",
        "    description='Design Mode:',\n",
        "    disabled=False,\n",
        "    style = {'description_width': 'initial'},\n",
        "    layout=common_layout,\n",
        ")\n",
        "\n",
        "# Define layout for right-aligning the description and common width for labels\n",
        "description_layout = widgets.Layout(display='flex', justify_content='flex-end', width='150px')\n",
        "\n",
        "# Define slider layout with common width\n",
        "slider_layout = widgets.Layout(width='250px')\n",
        "\n",
        "# Create sliders with the defined layout\n",
        "stiffness_slider = widgets.FloatSlider(\n",
        "    value=0.0,\n",
        "    min=0.0,\n",
        "    max=400.0,\n",
        "    step=0.1,\n",
        "    disabled=False,\n",
        "    continuous_update=False,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='.1f',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=slider_layout\n",
        ")\n",
        "\n",
        "rest_length_slider = widgets.FloatSlider(\n",
        "    value=0.0,\n",
        "    min=-1.0,\n",
        "    max=1.0,\n",
        "    step=0.01,\n",
        "    disabled=False,\n",
        "    continuous_update=False,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='.2f',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=slider_layout\n",
        ")\n",
        "\n",
        "# Create label widgets for alignment\n",
        "stiffness_label = widgets.Label(value='stiffness [N-m/rad]:', layout=description_layout)\n",
        "rest_length_label = widgets.Label(value='rest length [rad]:', layout=description_layout)\n",
        "\n",
        "# Combine labels and sliders into HBox\n",
        "stiffness_hbox = widgets.HBox([stiffness_label, stiffness_slider], layout=common_layout)\n",
        "rest_length_hbox = widgets.HBox([rest_length_label, rest_length_slider], layout=common_layout)\n",
        "\n",
        "# Create submit button with common layout\n",
        "submit_button = widgets.Button(\n",
        "    description='Submit',\n",
        "    disabled=False,\n",
        "    button_style='',\n",
        "    tooltip='Click to submit',\n",
        "    icon='check',\n",
        "    layout=common_layout\n",
        ")\n",
        "\n",
        "# Mode change handler\n",
        "def on_mode_change(change):\n",
        "    if change['new'] == 'Manual':\n",
        "        stiffness_slider.value = 200.0\n",
        "        rest_length_slider.value = 0.0\n",
        "        stiffness_slider.disabled = False\n",
        "        rest_length_slider.disabled = False\n",
        "    else:\n",
        "        stiffness_slider.value = 200.0\n",
        "        rest_length_slider.value = 0.0\n",
        "        stiffness_slider.disabled = True\n",
        "        rest_length_slider.disabled = True\n",
        "\n",
        "mode_selector.observe(on_mode_change, names='value')\n",
        "\n",
        "submit_button.on_click(on_submit)\n",
        "\n",
        "# Display widgets\n",
        "display(widgets.VBox([mode_selector, stiffness_hbox, rest_length_hbox, submit_button]))\n",
        "\n",
        "# Initialize state based on the current mode\n",
        "on_mode_change({'new': mode_selector.value})\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8.6. Modeling neural control with muscle synergies\n",
        "\n",
        "OpenSim now provides support for modeling muscle synergies via the new controller type [`SynergyController`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1SynergyController.html). This controller uses OpenSim [`Input`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1Input.html#a4d1e8b8a76452dc4d938742f19676fe3)s to pass in control signals from any component in the model that provides a scalar control signal via an [`Output`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1Output.html) (e.g., [`SignalGenerator`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1SignalGenerator.html)). The input control signals are multiplied by weighting vectors to generate a set of actuator controls. We refer to the input control signals as \"synergy excitations\" and the weighting vectors as \"synergy vectors\". In the following example, the synergies are used to produce muscle exctiations in the model, but `SynergyController` can be used with any OpenSim [`Actuator`](https://simtk.org/api_docs/opensim/api_docs451/classOpenSim_1_1Actuator.html), not just muscles.\n",
        "\n",
        "![muscle_synergies.png](https://drive.google.com/uc?export=download&id=14c3Ror0dh-beine2KHHyOpCLLD-GDUkc)\n",
        "\n",
        "Execute the cell below to load an interactive menu for running synergy-driven squat-to-stand motions. The menu provides a slider to select the number of muscle synergies between 1 and 8. As you decrease the number of synergies used by the `SynergyController`, the muscle control strategies will be become increasingly coupled. How will this affect the squat-to-stand motion?\n",
        "\n",
        "While running the synergy-driven optimizations, consider the following additional questions.\n",
        "\n",
        "**<u>Question 1.</u> What is the fewest number of synergies where the squat-to-stand motion is still possible?**\n",
        "\n",
        "**<u>Question 2.</u> Can you incorporate a passive device from the previous section to make standing possible with reduced neural control dimensionality?**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# @title Run a synergy-driven squat-to-stand optimization { display-mode: \"form\" }\n",
        "# @markdown Execute the cell to load the menu for specifying the number of synergies.\n",
        "\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display, Markdown\n",
        "import numpy as np\n",
        "\n",
        "# Suppress scikit-learn warnings.\n",
        "import warnings\n",
        "from sklearn.exceptions import ConvergenceWarning\n",
        "warnings.filterwarnings(\"ignore\", category=ConvergenceWarning)\n",
        "\n",
        "def printmd(string, color):\n",
        "    display(Markdown(f'<font color={color} size=\"4\">{string}</font>'))\n",
        "\n",
        "def run_synergy_driven_squat_to_stand(num_synergies):\n",
        "\n",
        "    # Part 1: Load the base squat-to-stand MocoStudy.\n",
        "    study = osim.MocoStudy(os.path.join('squat_to_stand', 'squat_to_stand.xml'))\n",
        "    problem = study.updProblem()\n",
        "    solver = osim.MocoCasADiSolver.safeDownCast(study.updSolver())\n",
        "\n",
        "    # Part 2: Recreate the model, now adding a passive assistive device at the knee\n",
        "    modelProcessor = osim.ModelProcessor('squatToStand_3dof9musc.osim')\n",
        "    modelProcessor.append(osim.ModOpScaleMaxIsometricForce(2.0))\n",
        "    modelProcessor.append(osim.ModOpReplaceMusclesWithDeGrooteFregly2016())\n",
        "    modelProcessor.append(osim.ModOpIgnorePassiveFiberForcesDGF())\n",
        "    modelProcessor.append(osim.ModOpIgnoreTendonCompliance())\n",
        "    modelProcessor.append(osim.ModOpIgnoreActivationDynamics())\n",
        "    modelProcessor.append(osim.ModOpScaleActiveFiberForceCurveWidthDGF(1.5))\n",
        "    modelProcessor.append(osim.ModOpReplacePathsWithFunctionBasedPaths(\n",
        "        os.path.join('path_fitting', 'squatToStand_FunctionBasedPathSet.xml')))\n",
        "    model = modelProcessor.process()\n",
        "    model.initSystem()\n",
        "\n",
        "    # Part 3: Load the solution to the squat-to-stand prediction to extract the muscle\n",
        "    # controls.\n",
        "    prevSolution = osim.MocoTrajectory(\n",
        "            os.path.join('squat_to_stand', 'squat_to_stand_prediction_solution.sto'))\n",
        "    controls = prevSolution.exportToControlsTable()\n",
        "\n",
        "    # Part 4: Use non-negative matrix factorization (NNMF) to extract a set of muscle\n",
        "    # synergies.\n",
        "    from sklearn.decomposition import NMF\n",
        "    nmf = NMF(n_components=num_synergies, init='random', random_state=0, max_iter=1000)\n",
        "\n",
        "    A = controls.getMatrix().to_numpy()\n",
        "    W = nmf.fit_transform(A)\n",
        "    H = nmf.components_\n",
        "\n",
        "    # Scale W and H assuming that the elements of H are all 0.5. This prevents the\n",
        "    # synergy excitations and synergy vector weights from being too small or too large.\n",
        "    scaleVec = 0.5*np.ones(H.shape[1])\n",
        "    for i in range(num_synergies):\n",
        "        scale = np.linalg.norm(scaleVec) / np.linalg.norm(H[i, :])\n",
        "        H[i, :] *= scale\n",
        "        W[:, i] /= scale\n",
        "\n",
        "    # Part 5: Add a SynergyController for the model's right leg.\n",
        "    controller = osim.SynergyController()\n",
        "    controller.setName(\"synergy_controller\")\n",
        "    # The number of actuators connected to the controller defines the number of\n",
        "    # weights in each synergy vector expected by the controller.\n",
        "    for name in controls.getColumnLabels():\n",
        "        controller.addActuator(\n",
        "                osim.Muscle.safeDownCast(model.getComponent(name)))\n",
        "    # Adding a synergy vector increases the number of synergies in the\n",
        "    # controller by one. This means that the number of Input control\n",
        "    # signals expected by the controller is also increased by one.\n",
        "    for i in range(num_synergies):\n",
        "        synergyVector = osim.Vector(H.shape[1], 0.0)\n",
        "        for j in range(H.shape[1]):\n",
        "            synergyVector.set(j, H[i, j])\n",
        "        controller.addSynergyVector(synergyVector)\n",
        "    model.addController(controller)\n",
        "    model.finalizeConnections()\n",
        "    model.initSystem()\n",
        "\n",
        "    # Part 6: Update the problem with the new model.\n",
        "    problem.setModel(model)\n",
        "\n",
        "    # Part 7: Set bounds on the synergy excitations and update the weight on the\n",
        "    # synergy excitations in the control effort cost term.\n",
        "    effort = osim.MocoControlGoal.safeDownCast(problem.updGoal(\"control_effort\"))\n",
        "    # Enabling this flag means muscle excitations are ignored in the effort cost.\n",
        "    effort.setIgnoreControlledActuators(True)\n",
        "    for i in range(num_synergies):\n",
        "        name = f'/controllerset/synergy_controller/synergy_excitation_{i}'\n",
        "        problem.setInputControlInfo(name, [0, 1.5*np.max(W)])\n",
        "        effort.setWeightForControl(name, 0.1)\n",
        "\n",
        "    # Part 8: Update the solver with the new problem.\n",
        "    solver.resetProblem(problem)\n",
        "\n",
        "    # Part 9: Solve!\n",
        "    solutionSealed = study.solve()\n",
        "    solution = solutionSealed.unseal()\n",
        "    if solution.success():\n",
        "        printmd(f'**Success!** Problem with {num_synergies} synergies converged!',\n",
        "                'green')\n",
        "    else:\n",
        "        printmd(f'**Failure.** Problem with {num_synergies} synergies did not converge.',\n",
        "                'red')\n",
        "\n",
        "    # This function computes the model control values from the\n",
        "    # SynergyController in the model and inserts them into the solution.\n",
        "    solution.generateControlsFromModelControllers(model)\n",
        "    # study.visualize(solution)\n",
        "\n",
        "    # Part 10: Write the solution to a Storage file.\n",
        "    if not os.path.exists('muscle_synergies'):\n",
        "            os.mkdir('muscle_synergies')\n",
        "    solutionFile = os.path.join('muscle_synergies',\n",
        "            f'squat_to_stand_{num_synergies}_muscle_synergies_solution.sto')\n",
        "    solution.write(solutionFile)\n",
        "\n",
        "    # Part 11: Create a PDF report of the solution.\n",
        "    output = os.path.join('muscle_synergies',\n",
        "            f'squat_to_stand_{num_synergies}_muscle_synergies_solution.pdf')\n",
        "    ref_files = []\n",
        "    ref_files.append(os.path.join('squat_to_stand',\n",
        "            'squat_to_stand_prediction_solution.sto'))\n",
        "    report = osim.report.Report(model,\n",
        "                                solutionFile,\n",
        "                                output=output,\n",
        "                                ref_files=ref_files,\n",
        "                                bilateral=True)\n",
        "    report.generate()\n",
        "\n",
        "def on_submit(b):\n",
        "    run_synergy_driven_squat_to_stand(slider.value)\n",
        "\n",
        "\n",
        "common_layout = widgets.Layout(width='300px')\n",
        "\n",
        "# Define layout for right-aligning the description and common width for labels\n",
        "description_layout = widgets.Layout(display='flex', justify_content='flex-end', width='100px')\n",
        "\n",
        "# Define slider layout with common width\n",
        "slider_layout = widgets.Layout(width='250px')\n",
        "\n",
        "slider = widgets.IntSlider(\n",
        "    value=5,\n",
        "    min=1,\n",
        "    max=8,\n",
        "    step=1,\n",
        "    disabled=False,\n",
        "    continuous_update=False,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    # readout_format='.1f',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=slider_layout\n",
        ")\n",
        "\n",
        "# Create label widgets for alignment\n",
        "slider_label = widgets.Label(value='no. synergies:', layout=description_layout)\n",
        "\n",
        "# Combine labels and sliders into HBox\n",
        "slider_hbox = widgets.HBox([slider_label, slider], layout=common_layout)\n",
        "\n",
        "submit_button = widgets.Button(\n",
        "    description='Submit',\n",
        "    disabled=False,\n",
        "    button_style='',\n",
        "    tooltip='Click to submit',\n",
        "    icon='check',\n",
        "    layout=common_layout\n",
        ")\n",
        "\n",
        "submit_button.on_click(on_submit)\n",
        "\n",
        "display(widgets.VBox([slider_hbox, submit_button]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J5mhS9zNX9ff"
      },
      "source": [
        "## 8.7. Conclusion"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_M04jPJLuiUz"
      },
      "source": [
        "In this tutorial you used [condacolab](https://github.com/conda-incubator/condacolab) to install [Miniconda](https://docs.conda.io/en/latest/miniconda.html) in [Google Colab](https://colab.research.google.com/?hl=en). Then, you installed Conda [OpenSim](https://opensim.stanford.edu/) with Moco into the Miniconda environment. You downloaded the 3 DOF, 9 muscle \"squat-to-stand\" model, fit function-based muscle paths, and used trajectory optimization to predict moving from a squat to a stand. Finally, you designed parameters for a passive assistive device at the knee and used muscle synergies to model how neural control limitations affect the squat-to-stand motion."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lLWgd8Noun6Y"
      },
      "source": [
        "## 8.8. Useful Links"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m8M3dlT5xQcD"
      },
      "source": [
        "> **OpenSim Website:** https://opensim.stanford.edu/\n",
        ">\n",
        "> **OpenSim Python Scripting:** https://opensimconfluence.atlassian.net/wiki/spaces/OpenSim/pages/53085346/Scripting+in+Python\n",
        ">\n",
        "> **OpenSim Moco Website:** https://opensim-org.github.io/opensim-moco-site/\n",
        ">\n",
        "> **OpenSim Moco Publication:** https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1008493\n",
        ">\n",
        "> **OpenSim SimTK Website:** https://simtk.org/projects/opensim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rFR2Rf9tk_9A"
      },
      "source": [
        "## 8.9. Acknowledgments\n",
        "\n",
        "Thanks to [OpenSimColab](https://simtk.org/projects/opencolab) project [[2]](https://doi.org/10.1080/10255842.2022.2104607) for creating the first OpenSim Conda package."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dxMjdukAnH2N"
      },
      "source": [
        "## 8.10. References"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WX2A_iaYulV1"
      },
      "source": [
        "> [1] Dembia CL, Bianco NA, Falisse A, Hicks JL, Delp SL (2020) **OpenSim Moco: Musculoskeletal optimal control.** PLOS Computational Biology 16(12): e1008493. https://doi.org/10.1371/journal.pcbi.1008493\n",
        ">\n",
        "> [2] Mokhtarzadeh, H., Jiang, F., Zhao, S., & Malekipour, F. (2022). **OpenColab project: OpenSim in Google colaboratory to explore biomechanics on the web.** *Computer Methods in Biomechanics and Biomedical Engineering*, 1–9. https://doi.org/10.1080/10255842.2022.2104607"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "opensim",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
